import cv2
import numpy as np
import time
import os
import ctypes
import threading

def load_yolo(use_tiny=True):
    """
    加载YOLO模型
    
    这个函数做了什么：
    1. 找到模型文件的路径
    2. 加载YOLO神经网络（可以选择标准版或轻量级tiny版本）
    3. 读取可以识别的物体类别名称
    
    就像是教会电脑认识各种物体（包括人）的眼睛
    
    参数：
        use_tiny: 是否使用轻量级的YOLOv3-tiny模型，速度更快但准确度略低
    """
    # 获取当前脚本所在的目录路径
    current_dir = os.path.dirname(os.path.abspath(__file__))
    
    # 根据选择设置模型文件路径
    if use_tiny:
        # 使用轻量级模型（速度更快，但准确度略低）
        config_path = os.path.join(current_dir, 'yolov3-tiny.cfg')  # 轻量级模型结构文件
        weights_path = os.path.join(current_dir, 'yolov3-tiny.weights')  # 轻量级模型权重文件
        
        # 如果轻量级模型文件不存在，尝试下载
        if not os.path.exists(config_path) or not os.path.exists(weights_path):
            print("轻量级模型文件不存在，将使用标准模型。如需更快速度，请下载yolov3-tiny.cfg和yolov3-tiny.weights文件")
            # 回退到标准模型
            config_path = os.path.join(current_dir, 'yolov3.cfg')
            weights_path = os.path.join(current_dir, 'yolov3.weights')
    else:
        # 使用标准模型（更准确，但速度较慢）
        config_path = os.path.join(current_dir, 'yolov3.cfg')
        weights_path = os.path.join(current_dir, 'yolov3.weights')
    
    classes_path = os.path.join(current_dir, 'coco.names')  # 类别名称文件
    
    # 检查文件是否存在
    if not os.path.exists(config_path) or not os.path.exists(weights_path) or not os.path.exists(classes_path):
        print(f"错误：模型文件不存在，请确保{os.path.basename(config_path)}、{os.path.basename(weights_path)}和coco.names文件在同一目录下")
        return None, None
    
    # 加载YOLO模型（就像是给电脑安装了一双能识别物体的眼睛）
    net = cv2.dnn.readNetFromDarknet(config_path, weights_path)
    
    # 尝试使用OpenCL加速（如果可用）
    cv2.ocl.setUseOpenCL(True)
    
    # 尝试使用CUDA（GPU）加速
    try:
        # 设置计算设备为GPU/CUDA
        net.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA)
        net.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA)
        print("成功启用CUDA GPU加速！")
    except Exception as e:
        print(f"无法使用CUDA GPU加速，将使用CPU: {e}")
        # 如果CUDA不可用，回退到CPU
        net.setPreferableBackend(cv2.dnn.DNN_BACKEND_OPENCV)
        net.setPreferableTarget(cv2.dnn.DNN_TARGET_CPU)
    
    # 读取类别名称（告诉电脑它能认出哪些物体，比如人、车、狗等）
    with open(classes_path, 'r') as f:
        classes = [line.strip() for line in f.readlines()]
    
    # 返回加载好的模型和类别名称
    return net, classes

def detect_people(frame, net, classes, confidence_threshold=0.4, nms_threshold=0.4, input_size=(320, 320), skip_frames=0):
    """
    检测图像中的人
    
    这个函数做了什么：
    1. 将图像转换为神经网络能理解的格式
    2. 让神经网络分析图像，找出所有物体
    3. 从所有物体中筛选出「人」
    4. 返回人的位置信息和数量
    
    就像是让电脑在照片中找出所有的人
    
    参数：
        frame: 要分析的图像帧
        net: 加载好的YOLO模型
        classes: 类别名称列表
        confidence_threshold: 置信度阈值，低于此值的检测结果会被忽略
        nms_threshold: 非极大值抑制阈值
        input_size: 输入神经网络的图像尺寸，越小速度越快，但准确度略低
        skip_frames: 是否跳过处理某些帧（0表示不跳过）
    """
    # 获取图像尺寸
    height, width, _ = frame.shape
    
    # 创建神经网络的输入（将图像转换为神经网络能理解的格式）
    # 这一步就像是把照片处理成特殊格式，让电脑的「眼睛」能看懂
    # 使用较小的输入尺寸可以加快处理速度
    blob = cv2.dnn.blobFromImage(frame, 1/255.0, input_size, swapRB=True, crop=False)
    
    # 设置网络输入
    net.setInput(blob)
    
    # 获取网络输出层的名称
    layer_names = net.getLayerNames()
    output_layers = [layer_names[i - 1] for i in net.getUnconnectedOutLayers()]
    
    # 前向传播（让神经网络分析图像）
    # 这一步就像是电脑在思考「照片里有什么物体」
    outputs = net.forward(output_layers)
    
    # 初始化检测结果列表
    boxes = []  # 存储检测到的人的位置（框）
    confidences = []  # 存储检测的置信度（电脑有多确定这是个人）
    class_ids = []  # 存储检测到的类别ID
    
    # 处理每个输出层
    for output in outputs:
        # 处理每个检测结果
        for detection in output:
            # 前5个值是框的信息和置信度，后面的值是每个类别的分数
            scores = detection[5:]
            # 找出得分最高的类别
            class_id = np.argmax(scores)
            confidence = scores[class_id]
            
            # 如果检测到的是人（class_id=0）且置信度大于阈值
            if class_id == 0 and confidence > confidence_threshold:  # 在COCO数据集中，人的类别ID是0
                # 计算边界框的坐标
                # detection[0:4]包含了边界框的中心坐标和宽高（相对值）
                center_x = int(detection[0] * width)
                center_y = int(detection[1] * height)
                w = int(detection[2] * width)
                h = int(detection[3] * height)
                
                # 计算左上角坐标
                x = int(center_x - w / 2)
                y = int(center_y - h / 2)
                
                # 保存结果
                boxes.append([x, y, w, h])
                confidences.append(float(confidence))
                class_ids.append(class_id)
    
    # 非极大值抑制（去除重复的检测框）
    # 这一步就像是去掉重复计算的人，确保每个人只被计数一次
    indices = cv2.dnn.NMSBoxes(boxes, confidences, confidence_threshold, nms_threshold)
    
    # 处理检测结果
    people_boxes = []
    if len(indices) > 0:  # 确保indices不为空
        for i in indices.flatten() if isinstance(indices, np.ndarray) else indices:
            # OpenCV 4.5.4之前的版本，indices是一个嵌套列表
            if isinstance(i, (list, tuple)):
                i = i[0]
            # 获取边界框信息
            box = boxes[i]
            people_boxes.append(box)
    
    # 返回检测到的人的边界框列表和人数
    return people_boxes, len(people_boxes)

def switch_to_desktop():
    """
    切换到桌面（相当于按Windows+D键）
    """
    try:
        # 使用Windows API模拟按下Windows+D键
        user32 = ctypes.windll.user32
        # VK_LWIN (0x5B) + VK_D (0x44)
        user32.keybd_event(0x5B, 0, 0, 0)  # Windows键按下
        user32.keybd_event(0x44, 0, 0, 0)  # D键按下
        user32.keybd_event(0x44, 0, 2, 0)  # D键释放
        user32.keybd_event(0x5B, 0, 2, 0)  # Windows键释放
        print("检测到两个人，已自动切换到桌面")
        
        # 为了避免过快切换回来，添加一个冷却时间
        time.sleep(3)  # 冷却时间3秒，避免频繁切换
    except Exception as e:
        print(f"切换到桌面时出错: {e}")

def check_cuda_availability():
    """
    检查是否支持CUDA（GPU加速）
    """
    # 检查OpenCV是否支持CUDA
    cv_info = cv2.getBuildInformation()
    cuda_support = "CUDA" in cv_info and "YES" in cv_info[cv_info.find("CUDA"):cv_info.find("\n", cv_info.find("CUDA"))]
    
    if cuda_support:
        print("检测到OpenCV支持CUDA GPU加速")
        return True
    else:
        print("警告：当前OpenCV版本不支持CUDA GPU加速，将使用CPU运算")
        print("提示：如需使用GPU加速，请安装支持CUDA的OpenCV版本")
        return False

def main():
    """
    主函数：启动摄像头，检测人体数量并实时显示
    """
    # 检查是否支持CUDA
    has_cuda = check_cuda_availability()
    
    # 加载YOLO模型和类别（使用轻量级模型以提高速度）
    net, classes = load_yolo(use_tiny=True)
    if net is None or classes is None:
        return
    
    # 打开摄像头（0表示默认摄像头）
    cap = cv2.VideoCapture(0)
    
    # 检查摄像头是否成功打开
    if not cap.isOpened():
        print("错误：无法打开摄像头")
        return
    
    # 设置摄像头分辨率（降低分辨率以提高速度）
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
    
    # 设置窗口
    cv2.namedWindow("人体计数器", cv2.WINDOW_NORMAL)
    cv2.resizeWindow("人体计数器", 800, 600)
    
    # 用于计算FPS的变量
    frame_count = 0
    start_time = time.time()
    fps = 0
    
    # 用于跳帧处理的变量（每隔几帧处理一次，提高速度）
    process_this_frame = 0
    skip_frames = 2  # 每3帧处理1帧
    
    # 保存上一次的检测结果
    last_people_boxes = []
    last_people_count = 0
    
    # 用于跟踪是否已经切换到桌面的标志
    desktop_switched = False
    last_switch_time = 0  # 上次切换到桌面的时间
    
    print("优化版人体计数器已启动。按'q'键退出。")
    if has_cuda:
        print("正在使用GPU加速的YOLO模型检测画面中的人...")
    else:
        print("正在使用CPU版YOLO模型检测画面中的人...")
    print("提示：如果仍然卡顿，可以按's'键增加跳帧数量以提高速度（但会降低响应速度）")
    print("特殊功能：当检测到两个人时，将自动切换到桌面")
    
    while True:
        # 读取一帧图像
        ret, frame = cap.read()
        
        # 如果读取失败，退出循环
        if not ret:
            print("错误：无法读取视频帧")
            break
        
        # 缩小图像尺寸以加快处理速度（可选）
        # frame = cv2.resize(frame, (640, 480))
        
        # 只处理部分帧以提高FPS
        if process_this_frame == 0:
            # 检测人体（使用较小的输入尺寸以提高速度）
            people_boxes, people_count = detect_people(frame, net, classes, 
                                                     input_size=(320, 320),  # 较小的输入尺寸
                                                     confidence_threshold=0.4)  # 稍微降低置信度阈值
            
            # 更新上一次的检测结果
            last_people_boxes = people_boxes
            last_people_count = people_count
            
            # 检查是否有两个人，如果有则切换到桌面
            current_time = time.time()
            if people_count == 2 and not desktop_switched and (current_time - last_switch_time) > 5:
                # 使用线程来切换到桌面，避免阻塞主线程
                threading.Thread(target=switch_to_desktop).start()
                desktop_switched = True
                last_switch_time = current_time
            elif people_count != 2:
                desktop_switched = False  # 重置标志，当人数不是2时
        else:
            # 使用上一次的检测结果
            people_boxes = last_people_boxes
            people_count = last_people_count
        
        # 更新处理帧计数器
        process_this_frame = (process_this_frame + 1) % (skip_frames + 1)
        
        # 在每个检测到的人体周围画一个矩形
        for (x, y, w, h) in people_boxes:
            cv2.rectangle(frame, (x, y), (x+w, y+h), (0, 255, 0), 2)
        
        # 计算FPS
        frame_count += 1
        elapsed_time = time.time() - start_time
        if elapsed_time > 1:  # 每秒更新一次FPS
            fps = frame_count / elapsed_time
            frame_count = 0
            start_time = time.time()
        
        # 在图像上显示人体数量和FPS
        cv2.putText(frame, f"人体数量: {people_count}", (10, 30), 
                    cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
        cv2.putText(frame, f"FPS: {fps:.1f}", (10, 60), 
                    cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
        
        # 显示结果
        cv2.imshow("人体计数器", frame)
        
        # 按'q'键退出，按's'键增加跳帧数
        key = cv2.waitKey(1) & 0xFF
        if key == ord('q'):
            break
        elif key == ord('s'):
            skip_frames += 1
            print(f"已增加跳帧数量，当前每{skip_frames + 1}帧处理1帧")
    
    # 释放资源
    cap.release()
    cv2.destroyAllWindows()
    print("人体计数器已关闭")

if __name__ == "__main__":
    main()
